import math
import torch
import numpy as np

class WeightScheduler:
    """
    Base class for weight schedulers.
    """
    def get_weight(self, step):
        raise NotImplementedError("Subclasses should implement this method.")

class CosDown(WeightScheduler):
    """
    Cosine Down Scheduler for loss weight adjustment.
    """
    def __init__(
            self,
            total_steps,
            init_weight=1.0,
            min_weight=0
        ):
        self.tot_steps = total_steps
        self.init_weight = init_weight
        self.min_weight = min_weight

    def get_weight(self, step):
        """
        Get the current weight based on the cosine schedule.
        :param step: Current training step.
        :return: Adjusted weight for the loss function.
        """
        if step >= self.tot_steps:
            return self.min_weight
        else:
            cos_value = (1 + math.cos(math.pi * step / self.tot_steps)) / 2
            return self.min_weight + (self.init_weight - self.min_weight) * cos_value

# TESTING
if __name__ == "__main__":
    scheduler = CosDown(total_steps=1000, init_weight=1.0, min_weight=0.0)
    for step in range(0, 1001, 10):
        weight = scheduler.get_weight(step)
        print(f"Step {step}: Weight = {weight:.4f}")